# routes.py
from flask import request, jsonify, render_template, url_for, redirect, session, Blueprint
from models import Contact, db, Email, MailRecipient
from datetime import datetime
from ai_response import AIResponse
from werkzeug.security import generate_password_hash, check_password_hash
import imaplib
import email
from email.header import decode_header

main_bp = Blueprint('main', __name__)
ai_response = AIResponse()

# 主页面
@main_bp.route('/')
def index():
    return render_template('index.html')

@main_bp.route('/register', methods=['GET', 'POST'])
def register():
    if request.method == 'POST':
        username = request.form['username']
        password = request.form['password']
        email = request.form['email']
        cellphone = request.form['cellphone']
        address = request.form['address']

        if Contact.query.filter_by(username=username).first() is not None:
            return jsonify({'message': '用户名已存在'}), 400
        if Contact.query.filter_by(email=email).first() is not None:
            return jsonify({'message': '电子邮箱已注册'}), 400

        hashed_password = generate_password_hash(password, method='pbkdf2:sha256')
        new_user = Contact(username=username, password=hashed_password, email=email, cellphone=cellphone, address=address)
        db.session.add(new_user)
        db.session.commit()
        return jsonify({'message': '注册成功'}), 200
    else:
        return render_template('register.html')


# 用户登录
@main_bp.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        username = request.form['username']
        password = request.form['password']
        # 查询 User 表以验证用户名和密码
        user = Contact.query.filter_by(username=username).first()
        if (user and check_password_hash(user.password, password)) or (user and user.password == password):
            # 用户名和密码验证成功后，查询 Contact 表以获取 uid
            contact = Contact.query.filter_by(username=username).first()
            if contact:
                # 将用户的 uid 存储在 session 中以便后续使用
                session['user_id'] = contact.uid
                return redirect(url_for('main.dashboard'))
            else:
                return redirect(url_for('main.login', error='找不到对应的联系信息'))
        else:
            return redirect(url_for('main.login', error='用户名或密码错误'))
    else:
        error = request.args.get('error')
        return render_template('login.html', error=error)

@main_bp.route('/dashboard')
def dashboard():
    # 确保用户已登录
    if 'user_id' not in session:
        return redirect(url_for('main.login'))
    return render_template('dashboard.html')


# 发送邮件
@main_bp.route('/send_email', methods=['GET', 'POST'])
def send_email():
    if request.method == 'POST':
        current_user_id = session.get('user_id')
        if not current_user_id:
            return jsonify({'status': 'error', 'message': '用户未登录'}), 401

        sender_contact = Contact.query.filter_by(uid=current_user_id).first()
        if not sender_contact:
            return jsonify({'status': 'error', 'message': '找不到发件人的联系信息'}), 404

        receiver_email = request.form['receiver']
        subject = request.form['subject']
        body = request.form['body']

        receiver_contact = Contact.query.filter_by(email=receiver_email).first()
        if not receiver_contact:
            return jsonify({'status': 'error', 'message': '找不到收件人的联系信息'}), 404

        new_email = Email(title=subject, uid=current_user_id, textbody=body)
        db.session.add(new_email)
        db.session.flush()

        new_recipient = MailRecipient(eid=new_email.eid, uid=receiver_contact.uid)
        db.session.add(new_recipient)
        db.session.commit()

        return jsonify({'status': 'success'}), 200
    else:
        success = request.args.get('success')
        return render_template('send_email.html', success=success)
    
#查看收件箱
@main_bp.route('/inbox')
def inbox():
    current_user_id = session.get('user_id')
    if not current_user_id:
        return jsonify({'message': '用户未登录'}), 401

    # 查询当前用户未删除的邮件，包含发送人信息
    received_emails = (
        Email.query
        .join(MailRecipient, MailRecipient.eid == Email.eid)
        .filter(MailRecipient.uid == current_user_id, Email.is_deleted == False)
        .all()
    )

    email_senders = []
    for email in received_emails:
        sender = Contact.query.get(email.uid)
        email_info = {
            'eid': email.eid,
            'title': email.title,
            'create_time': email.create_time,
            'sender_id': sender.uid 
        }
        email_senders.append(email_info)

    return render_template('inbox.html', emails=email_senders)

#查看接收的邮件内容
@main_bp.route('/view_email/<int:eid>')
def view_email(eid):
    # 确保用户已登录
    if 'user_id' not in session:
        return jsonify({'message': '用户未登录'}), 401

    # 根据邮件ID查询邮件内容
    email = Email.query.get_or_404(eid)
    return render_template('view_email.html', email=email)

#查看已发送邮件内容
@main_bp.route('/view_email2/<int:eid>')
def view_email2(eid):
    # 确保用户已登录
    if 'user_id' not in session:
        return jsonify({'message': '用户未登录'}), 401

    # 根据邮件ID查询邮件内容
    email = Email.query.get_or_404(eid)
    return render_template('view_email2.html', email=email)

#查看已发送邮件
@main_bp.route('/sent')
def sent_emails():
    current_user_id = session.get('user_id')
    if not current_user_id:
        return jsonify({'message': '用户未登录'}), 401

    # 查询当前用户发送的所有邮件
    sent_emails = Email.query.filter_by(uid=current_user_id, is_deleted=False).all()

    # 创建包含邮件和接收人信息的列表
    email_recipients = []
    for email in sent_emails:
        recipients = MailRecipient.query.filter_by(eid=email.eid).all()
        recipient_ids = [recipient.uid for recipient in recipients]
        email_info = {
            'eid': email.eid,
            'title': email.title,
            'create_time': email.create_time,
            'recipient_ids': recipient_ids
        }
        email_recipients.append(email_info)

    return render_template('sent_emails.html', emails=email_recipients)

#编辑邮件
@main_bp.route('/edit_email/<int:eid>', methods=['GET', 'POST'])
def edit_email(eid):
    # 确保用户已登录
    if 'user_id' not in session:
        return jsonify({'message': '用户未登录'}), 401

    email = Email.query.get_or_404(eid)
    if request.method == 'POST':
        email.title = request.form['title']
        email.textbody = request.form['body']
        email.create_time = datetime.utcnow()
        db.session.commit()
        return redirect(url_for('main.sent_emails'))
    return render_template('edit_email.html', email=email)

#删除邮件
@main_bp.route('/delete_email/<int:eid>', methods=['POST'])
def delete_email(eid):
    current_user_id = session.get('user_id')
    if not current_user_id:
        return jsonify({'message': '用户未登录'}), 401

    # 确认当前用户是否为邮件接收人
    recipient = MailRecipient.query.filter_by(eid=eid, uid=current_user_id).first()
    if not recipient:
        return jsonify({'message': '无权删除此邮件'}), 403

    email = Email.query.get_or_404(eid)
    email.is_deleted = True
    db.session.commit()
    return redirect(url_for('main.inbox'))

# 查看已删除邮件
@main_bp.route('/deleted_emails')
def deleted_emails():
    current_user_id = session.get('user_id')
    if not current_user_id:
        return jsonify({'message': '用户未登录'}), 401

    # 查询当前用户标记为已删除的邮件
    deleted_emails = (
        Email.query
        .join(MailRecipient, MailRecipient.eid == Email.eid)
        .filter(MailRecipient.uid == current_user_id, Email.is_deleted == True)
        .all()
    )
    return render_template('deleted_emails.html', emails=deleted_emails)

# 恢复已删除邮件
@main_bp.route('/restore_email/<int:eid>', methods=['POST'])
def restore_email(eid):
    # ... 省略验证用户登录的代码 ...
    email = Email.query.get_or_404(eid)
    if email.uid != session['user_id']:
        return jsonify({'message': '无权恢复此邮件'}), 403

    email.is_deleted = False
    db.session.commit()
    return redirect(url_for('main.deleted_emails'))

# 永久删除邮件
@main_bp.route('/permanent_delete_email/<int:eid>', methods=['POST'])
def permanent_delete_email(eid):
    # 确保用户已登录
    if 'user_id' not in session:
        return jsonify({'message': '用户未登录'}), 401
    email = Email.query.get_or_404(eid)
    
    # 查询并删除所有关联的MailRecipient记录
    mail_recipients = MailRecipient.query.filter_by(eid=eid).all()
    for recipient in mail_recipients:
        db.session.delete(recipient)

    # 从数据库中永久删除邮件
    db.session.delete(email)
    db.session.commit()
    
    return redirect(url_for('main.deleted_emails'))


#自动回复
@main_bp.route('/ai_reply/<int:eid>', methods=['GET'])
def ai_reply(eid):
    if 'user_id' not in session:
        return redirect(url_for('main.login'))
    email = Email.query.get_or_404(eid)
    ai_reply_content = ai_response.get_response(email.textbody)
    return render_template('ai_reply.html', original_email=email, ai_reply=ai_reply_content)

#发送自动回复邮件
@main_bp.route('/send_ai_reply/<int:eid>', methods=['POST'])
def send_ai_reply(eid):
    if 'user_id' not in session:
        return redirect(url_for('main.login'))
    original_email = Email.query.get_or_404(eid)
    ai_reply = request.form['ai_reply']
    
    # 创建并保存AI回复邮件
    current_user_id = session.get('user_id')
    new_email = Email(title=f"Re: {original_email.title}", uid=current_user_id, textbody=ai_reply)
    db.session.add(new_email)
    db.session.flush()
    new_recipient = MailRecipient(eid=new_email.eid, uid=original_email.uid)
    db.session.add(new_recipient)
    db.session.commit()

    return redirect(url_for('main.inbox'))

# AI辅助生成
@main_bp.route('/ai_generate_email', methods=['POST'])
def ai_generate_email():
    data = request.get_json()
    prompt = data.get('prompt')
    
    # 使用AI模型生成邮件内容
    ai_reply_content = ai_response.assist_write(prompt)
    return jsonify({'emailContent': ai_reply_content})

# 导入其他邮箱邮件
@main_bp.route('/import_email', methods=['GET', 'POST'])
def import_email():
    if request.method == 'POST':
        email_address = request.form['email']
        password = request.form['password']
        
        try:
            # 使用 with 语句确保资源被正确释放
            with imaplib.IMAP4_SSL("imap.qq.com") as mail:
                mail.login(email_address, password)
                mail.select("inbox")
                
                # 获取所有邮件ID
                status, messages = mail.search(None, "ALL")
                mail_ids = messages[0].split()
                
                emails = fetch_emails(mail, mail_ids)  # 拆分邮件获取逻辑到子函数

                return render_template('view_imported_emails.html', emails=emails)

        except imaplib.IMAP4.error as e:  # 指定捕获的异常类型
            return jsonify({"message": "邮箱验证失败，错误信息：" + str(e)})

    return render_template('import_email.html')

def decode_text(text, encodings):
    for encoding in encodings:
        try:
            return text.decode(encoding)
        except (UnicodeDecodeError, TypeError):
            continue
    return text.decode("utf-8", errors="ignore")  # 默认回退到UTF-8并忽略错误

def fetch_emails(mail, mail_ids):
    encodings = ["utf-8", "gbk", "iso-8859-1", "ascii"]
    emails = []
    for email_id in mail_ids:
        # 获取邮件内容
        status, msg_data = mail.fetch(email_id, "(RFC822)")
        raw_email = msg_data[0][1]
        
        # 尝试逐行解码
        email_message = email.message_from_bytes(raw_email)
        
        # 提取邮件内容
        subject, encoding = decode_header(email_message["Subject"])[0]
        if isinstance(subject, bytes):
            subject = decode_text(subject, encodings)
                
        from_ = email_message.get("From")
        body = ""
        
        if email_message.is_multipart():
            for part in email_message.walk():
                content_type = part.get_content_type()
                if "attachment" not in part.get("Content-Disposition", ""):
                    payload = part.get_payload(decode=True)
                    if payload:
                        body = payload.decode("utf-8", errors="ignore")
                        break
        else:
            payload = email_message.get_payload(decode=True)
            if payload:
                body = payload.decode("utf-8", errors="ignore")
        
        emails.append({"from": from_, "subject": subject, "body": body})
    return emails

def save_emails_to_db(emails):
    for email_data in emails:
        # 获取发送者邮箱地址和主题
        from_email = email_data['from']
        subject = email_data['subject']
        body = email_data['body']
        
        # 假设你有一个当前用户的 ID
        current_user_id = session.get('user_id')
        
        # 检查发送者是否在联系人表中
        sender_contact = Contact.query.filter_by(email=from_email).first()
        
        if not sender_contact:
            # 如果发送者不在联系人表中，则添加
            sender_contact = Contact(
                username=from_email.split('@')[0],  # 使用邮箱用户名作为联系人用户名
                email=from_email,
                password='',  # 或者你可以生成一个随机密码
                cellphone='',
                address=''
            )
            db.session.add(sender_contact)
            db.session.flush()  # 获取新的 UID
        
        # 添加邮件到 emails 表
        new_email = Email(
            title=subject,
            uid=current_user_id,
            textbody=body
        )
        db.session.add(new_email)
        db.session.flush()  # 获取新的 EID
        
        # 添加邮件收件人到 mail_recipients 表
        new_recipient = MailRecipient(
            eid=new_email.eid,
            uid=sender_contact.uid
        )
        db.session.add(new_recipient)
    
    db.session.commit()
